import tensorflow as tf
from tensorflow.keras import layers, optimizers, datasets, Sequential
import os
from model import resnet18

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
tf.random.set_seed(2345)


def preprocess(x, y):
    x = tf.cast(x, dtype=tf.float32) / 256
    y = tf.cast(y, dtype=tf.int32)
    return x, y


# (x, y), (x_test, y_test) = datasets.cifar100.load_data()
# print(x.shape, y.shape, x_test.shape, y_test.shape)
model = resnet18()
model.build(input_shape=[None, 32, 32, 3])
model.summary()

# train_db = tf.data.Dataset.from_tensor_slices((x, y))
# train_db = train_db.shuffle(1000).map(preprocess).batch(64)

# test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
# test_db = test_db.map(preprocess).batch(64)

# sample = next(iter(train_db))


def main():
    model = resnet18()
    model.build(input_shape=[None, 32, 32, 3])
    model.summary()
    optimizer = optimizers.Adam(lr=1e-3)

    for epoch in range(50):
        for step, (x, y) in enumerate(train_db):

            with tf.GradientTape() as tape:

                logits = model(x)
                # [b] =>[b, 100]
                y_onehot = tf.one_hot(y, depth=100)
                # 计算损失
                loss = tf.losses.categorical_crossentropy(y_onehot, logits, from_logits=True)
                loss = tf.reduce_mean(loss)
            grads = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(grads, model.trainable_variables)

            if step % 100 == 0:
                print(epoch, step, 'loss: ', float(loss))

        total_num = 0
        total_correct = 0

        for x, y in test_db:
            logits = model(x)
            prod = tf.nn.softmax(logits, axis=1)
            pred = tf.argmax(prod, axis=1)
            pred = tf.cast(pred, dtype=tf.int32)

            correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
            correct = tf.reduce_sum(correct)
            total_num += x.shape[0]
            total_correct += int(correct)
        acc = total_correct / total_num
        print("acc_test: ", acc)


# if __name__ == "__main__":
#     main()
